package e_union_find.c_union_find_opt_by_size;

/**
 * 优化： 将定向合并改成选择性合并 （始终将小组向大组合并，将大大提高find的性能）
 * <p>
 * 之前存在的问题：
 * 1.虽然union的速度变快了，但是查找的速度也下降了
 * 2.之前一直采用的定向合并的方式（q向p合并）可能导致形成的树过高 大大降低find效率
 * <p>
 * 优化：
 * 1.维护每个组的大小sz[]
 * 2.合并的时候判断p和q的组容量，始终将小组向大组合并，将大大提高find的性能
 * <p>
 * 存在的问题：
 * 虽然基于size的优化很容易想到
 * 但是仍然存在一些情况导致小组向大组合并反而组的高度增加了
 * <p>
 * 继续优化的思路：始终从层次矮的组向层次高的组合并 - 保证树的高度不会继续增长
 */
public class UnionFind3 {
    private int[] pids; // 为了方便， 索引代表元素id，值代表该位置的元素所指向的pid
    private int[] size; // size[i]代表以i为根的组的成员数
    private int count; // 并查集包含的元素个数

    /**
     * 传入 size
     * 构造一个以size为大小的并查集，组号从0递增到size-1
     *
     * @param size
     */
    public UnionFind3(int size) {
        this.count = size;
        this.pids = new int[size];
        this.size = new int[size]; // 组成员人数不包括根，因初始值设为0
        for (int i = 0; i < size; i++) {
            this.pids[i] = i;
        }
    }

    /**
     * 查找元素所在组的根id - O(1~n)
     *
     * @param p 待查找元素的id
     * @return 带查找元素所在组的根id
     */
    public int find(int p) {
        assert p >= 0 && p < this.count;
        // 向上找到组的根id
        while (p != pids[p]) {
            p = pids[p];
        }
        return p;
    }

    /**
     * 判断两个元素id之间是否处于相同的组  - O(1~n)
     *
     * @param p
     * @param q
     * @return
     */
    public boolean isConnected(int p, int q) {
        return find(p) == find(q);
    }

    /**
     * 将p和q所在组联合成一个组  - O(1~n)
     * 将q所在的组的所有成员都拉入p所在的组中
     * 即让q所在组的根指向p所在组的根即可
     *
     * @param p
     * @param q
     */
    public void union(int p, int q) {
        assert (p >= 0 && p < this.count) && (q >= 0 && q < this.count);

        int pRootID = find(p);
        int qRootID = find(q);

        if (pRootID == qRootID) return;
        if (size[pRootID] < size[qRootID]) { // q组大
            pids[pRootID] = qRootID; // 合并到q组
            size[qRootID] += size[pRootID] + 1; // 1是 另一组的 根成员
        } else { // q组小
            pids[qRootID] = pRootID; // 合并到p组
            size[pRootID] += size[qRootID] + 1; // 1是 另一组的 根成员
        }
    }

    /**
     * 将所有元素都归入第一个元素的组中 - O(n)
     *
     * @param oids
     */
    public void unionAll(int... oids) {
        assert oids != null && oids.length > 1;
        int len = oids.length;
        for (int i = 1; i < len; i++) {
            union(oids[i-1],oids[i]);
        }
    }

    // TODO test method .. to be deleted
    public void testPrint() {
        System.out.print("id  -> ");
        for (int i = 0; i < count; i++) {
            System.out.print("\t" + i);
        }
        System.out.println();
        System.out.print("tid -> ");
        for (int i = 0; i < count; i++) {
            System.out.print("\t" + find(pids[i]));
        }
        System.out.println();
        System.out.print("tsz -> ");
        for (int i = 0; i < count; i++) {
            System.out.print("\t" + size[i]);
        }
        System.out.println("\n=======================");
    }

}
